import os
import pandas as pd
import copy
import json

from utlis import use_llm_api, extract_nested_dic_vals, flatten_list
from kb_encoder import replace_items


def check_common_errs(filename, KB_PATH, llm_apis, model_config, USER_SETTINGS, api_name='qwen_api', llm_histories=[]):
    try:
        file_path = os.path.join(KB_PATH, filename, 'KB_PTXT.csv')
        file_df = pd.read_csv(file_path, encoding='utf-8', usecols=['path', 'content', 'linkage'])
    except Exception as e:
        print(e)
        print('check the file path, the file may not be parsed yet...')
        
    identified_errs = []
    modified_contents = []
    failures = []
    
    for i, row in file_df.iterrows():
        content = row['content']
        root = row['path'].split('-->')[0]
        
        record_dic = json.loads(content)
        contents_ = extract_nested_dic_vals(record_dic)
        
        replaced_contents = []
        replaced_labels = []
        for content in contents_:
            temp_contents, temp_labels = replace_items(root, content, KB_PATH, mode='split')
            replaced_contents.append(temp_contents)
            replaced_labels.append(temp_labels)
        
        labels_ = flatten_list(replaced_labels)
        contents_ = flatten_list(replaced_contents)
        
        content4inspect = '\n'.join(contents_)
        original_contents = copy.deepcopy(content4inspect)     
        response, llm_histories = use_llm_api(llm_apis[api_name],
                                                histories=llm_histories,
                                                paras={ 'task':'inspect-typos', 
                                                        'query' : '',
                                                        'texts':content4inspect, 
                                                        'local_model_name':'qwen',
                                                        'model':'qwen-plus' },
                                                config=model_config,
                                                settings=USER_SETTINGS,
                                                record_his=False)

        try:
            corrections = response['回答']
            if len(corrections)>0:
                identified_errs.extend(corrections)
                for ct in corrections:
                    content4inspect = content4inspect.replace(ct['原词'], ct['修正'])
        except Exception as e:
            print('\t', e)
            print(response)
            failures.append(content4inspect)
        modified_contents.append(content4inspect)
        
    identified_errs = pd.DataFrame(identified_errs)
    return identified_errs, '\n'.join(modified_contents), original_contents, failures

        
if __name__ =="__main__":
    from META import USER_SETTINGS, KB_PATH, model_config, llm_apis
    
    file_name = 'D:\Local codes\知识固化库_DEMO\Supplementary Files\安全大模型技术与市场研究报告.docx'
    file_path = os.path.join(KB_PATH, file_name.split(os.sep)[-1], 'KB_PTXT.csv')
    
    identified_errs, modified_contents, original_contents, failures = check_common_errs(file_path, KB_PATH, llm_apis)
            
    print(identified_errs)
    
    
    